Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss
面向 LLM 的高效知识蒸馏：离线 Top-K logit 缓存与融合分块 KL 损失
🧭 训练效率 · 知识蒸馏 · 长上下文 ★ 深度解读
Multiverse Computing
通过缓存教师 Top-K logits 和融合分块 KL 损失，实现单 GPU 上长上下文蒸馏训练，匹配在线质量且内存线性增长。
📌 工作参考价值（4 点）
- 离线蒸馏策略：预教师 Top-K logit 一次缓存后重复用于多轮消融实验，大幅降低资源占用，适合需频繁蒸馏的推荐模型迭代。
- 长上下文训练解锁：融合分块 KL 损失让峰值内存与序列长度成线性关系，可在单卡 H200 上训练 32K 上下文，用于处理电商搜索中长用户行为序列的场景。
- 损失设计经验：logit KL 损失是知识迁移的必需品，追加隐藏状态特征损失（MSE 或余弦）可稳定提升约 1% MMLU。
- 打包训练简化：蒸馏时用朴素的全 1 注意力掩码做序列打包，仅牺牲约 1 个 MMLU 点，工程实现简单且效率高。